import zlib, bz2, os, math

from strange_io.text import StrangeText
from strange_io import SIGNATURE, get_format, _load_bytearray
import strange_io.dynamic_integer as dynamic_integer

def none_compression(data):
    return data

class CrunchArchive:
    def __init__(self, name='noname', type='none'):
        self.name = name
        self.type = type
        self.files = {}
        self.dirs = []
        self.compressed = False
        self.methods = {
            'none': {
                'compress': none_compression,
                'decompress': none_compression
            },
            'zlib': {
                'compress': zlib.compress,
                'decompress': zlib.decompress
            },
            'bz2': {
                'compress': bz2.compress,
                'decompress': bz2.decompress
            }
        }

    def dump(self) -> dict:
        out = {}
        out['name'] = self.name
        out['type'] = self.type
        out['files'] = self.files
        out['dirs'] = self.dirs
        out['compressed'] = self.compressed
        out['methods'] = self.methods
        return out

    def load(self, archive: dict):
        self.name = archive['name']
        self.type = archive['type']
        self.files = archive['files']
        self.dirs = archive['dirs']
        self.compressed = archive['compressed']
        if 'methods' in archive:
            self.methods = archive['methods']

    def dump_file(self, path):
        out = b''
        header = SIGNATURE + b'crunch\n'
        out += header
        out += self.dump_bytes()
        self.compress()
        with open(path, 'wb') as f:
            f.write(out)
            f.close()

    def load_file(self, path, ignore_type=False) -> dict:
        if get_format(path) == b'crunch' or ignore_type:
            file = _load_bytearray(path, remove_type=True)
            self.load_bytes(file)
        else:
            raise TypeError('Attempted to load non-crunch archive document as crunch archive')

    def dump_bytes(self) -> bytes:
        archive = self.dump()
        out = self.encode(archive)
        return out

    def load_bytes(self, archive: bytes):
        out = self.decode(archive)
        self.load(out)

    def create_dir(self, path):
        if not self.path_exists(path):
            self.dirs.append(path)
        else:
            raise FileExistsError()

    def create_file(self, path, method='zlib', data=b''):
        if not self.path_exists(path):
            if self.compressed:
                data = self.methods[method]['compress'](data)
            self.files[path] = {
                'data': data,
                'method': method
            }
        else:
            raise FileExistsError("A file with that path already exists")

    def add_file(self, system_path, archive_path, method='zlib'):
        if os.path.exists(system_path):
            if os.path.isfile(system_path):
                if not self.path_exists(archive_path):
                    with open(system_path, "rb") as f:
                        self.create_file(archive_path, method, f.read())
                else:
                    raise FileExistsError("Archive path already exists")
            else:
                raise IsADirectoryError("System path leads to a directory")
        else:
            raise FileNotFoundError("System path doesn't exist")

    def path_exists(self, path):
        return path in self.dirs or path in self.files

    def get_file(self, path:str, compressed:bool|None=None) -> tuple[bytes, bool]:
        if compressed is None:
            return self.files[path], self.compressed
        elif compressed is bool:
            if compressed:
                self.compress()
                return self.files[path], self.compressed
            else:
                self.decompress()
                return self.files[path], self.compressed
        else:
            return TypeError()

    def compress(self):
        if not self.compressed:
            for file in self.files:
                method = self.files[file]['method']
                method_func = self.methods[method]['compress']
                self.files[file]['data'] = method_func(self.files[file]['data'])
            self.compressed = True

    def decompress(self):
        if self.compressed:
            for file in self.files:
                method = self.files[file]['method']
                method_func = self.methods[method]['decompress']
                self.files[file]['data'] = method_func(self.files[file]['data'])
            self.compressed = False

    @staticmethod
    def encode(archive: dict) -> bytes:
        out = b''
        text = StrangeText()
        #region name
        number = len(archive['name'])
        out += dynamic_integer.encode(number)[0]
        text.load(archive['name'])
        out += text.dump_bytes()
        #endregion
        #region type
        number = len(archive['type'])
        out += dynamic_integer.encode(number)[0]
        text.load(archive['type'])
        out += text.dump_bytes()
        #endregion
        #region number of files
        number = len(archive['files'])
        out += dynamic_integer.encode(number)[0]
        #endregion
        for file in archive['files']:
            #region file path
            number = len(file)
            out += dynamic_integer.encode(number)[0]
            text.load(file)
            out += text.dump_bytes()
            #endregion
            #region compression method
            number = len(archive['files'][file]['method'])
            out += dynamic_integer.encode(number)[0]
            text.load(archive['files'][file]['method'])
            out += text.dump_bytes()
            #endregion
            #region file
            number = len(archive['files'][file]['data'])
            out += dynamic_integer.encode(number)[0]
            out += archive['files'][file]['data']
            #endregion
        #region number of dirs
        number = len(archive['dirs'])
        out += dynamic_integer.encode(number)[0]
        #endregion
        for dir in archive['dirs']:
            number = len(archive['name'])
            out += dynamic_integer.encode(number)[0]
            text.load(dir)
            out += text.dump_bytes()
            del number, length
        return out

    @staticmethod
    def decode(archive: bytes) -> dict:
        file = bytearray(archive)
        text = StrangeText()
        out = {}
        out['compressed'] = True
        #region archive name
        name_length = dynamic_integer.decode(file)[0]
        archive_name = bytearray()
        for i in range(name_length):
            if i >= name_length:
                break
            archive_name += bytearray([file[i]])
            i += 1
        for i in range(name_length):
            del file[0]
        text.load_bytes(archive_name)
        out['name'] = text.dump()
        #endregion
        #region archive type
        type_length = dynamic_integer.decode(file)[0]
        archive_type = bytearray()
        for i in range(type_length):
            if i >= type_length:
                break
            archive_type.append(file[i])
            i += 1
        for i in range(type_length):
            del file[0]
        text.load_bytes(archive_type)
        out['type'] = text.dump()
        #endregion
        file_quantity = dynamic_integer.decode(file)[0]
        #region files
        out['files'] = {}
        for _ in range(file_quantity):
            #region file_path
            file_path_length = dynamic_integer.decode(file)[0]
            _bytes = bytearray()
            for _ in range(file_path_length):
                _bytes.append(file[0])
                del file[0]
            text.load_bytes(_bytes)
            file_path = text.dump()
            out['files'][file_path] = {}
            #endregion
            #region compression method
            file_compression_method_length = dynamic_integer.decode(file)[0]
            _bytes = bytearray()
            for _ in range(file_compression_method_length):
                _bytes.append(file[0])
                del file[0]
            text.load_bytes(_bytes)
            method = text.dump()
            out['files'][file_path]['method'] = method
            #endregion
            #region file
            file_size = dynamic_integer.decode(file)[0]
            _bytes = bytearray()
            for i in range(file_size):
                _bytes.append(file[0])
                del file[0]
            file_data = _bytes
            out['files'][file_path]['data'] = file_data
            #endregion
        #endregion
        dir_quantity = dynamic_integer.decode(file)[0]
        #region dirs
        out['dirs'] = []
        for _ in range(dir_quantity):
            #region dir_path
            dir_path = dynamic_integer.decode(file)[0]
            _bytes = bytearray()
            for i in range(dir_path_length):
                _bytes.append(file[i])
            for _ in range(dir_path_length):
                del file[0]
            text.load_bytes(_bytes)
            dir_path = text.dump()
            out['dirs'].append(dir_path)
            #endregion
        #endregion
        return out
